package main

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"rpc/msg"
	"strconv"
	"sync"
	"time"
)

var (
	pool = &sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 10240)
		}}

	client8bPool = &sync.Pool{
		New: func() interface{} {
			return make([]byte, 8)
		}}
)

func main() {

	var wg sync.WaitGroup

	client := NewClient(8, "127.0.0.1:8080")

	now := time.Now()
	for i := 0; i < 1000000; i++ {
		wg.Add(1)
		go func() {
			_, _ = client.Send([]byte("this is 2015 year" + strconv.Itoa(i)))
			//fmt.Println(string(v))
			wg.Done()
		}()
	}

	wg.Wait()
	fmt.Println("elapse: ", time.Since(now))

}

const (
	BatchSize  = 1 * 1024 * 1024
	BatchSize2 = 10000
)

type Client struct {
	connection []*Conn
	ack        *sync.Map
}

type Conn struct {
	con         net.Conn
	lock        *sync.Mutex
	state       int
	waitSendMsg chan *msg.Msg
}

func NewClient(connPool int, addr string) *Client {
	cli := Client{
		connection: make([]*Conn, 0),
		ack:        new(sync.Map),
	}

	for i := 0; i < connPool; i++ {
		conn, err := net.Dial("tcp4", addr)
		if err != nil {
			panic(err)
		}
		cli.connection = append(cli.connection, &Conn{
			con:         conn,
			lock:        new(sync.Mutex),
			state:       0,
			waitSendMsg: make(chan *msg.Msg, 102400),
		})
	}

	cli.doRead()
	cli.doSend2()

	return &cli
}

func (client *Client) Send(data []byte) ([]byte, error) {

	message := msg.NewMsg(data)
	resultChan := make(chan *msg.Msg)
	client.ack.Store(message.SeqNo, resultChan)
	defer func() {
		client.ack.Delete(message.SeqNo)
	}()

	client.asyncSend(message)

	rep := <-resultChan
	return rep.Data, nil
}

func (client *Client) asyncSend(message *msg.Msg) {
	i := time.Now().UnixMicro() & (int64(len(client.connection)) - 1)
	waitSendMsg := client.connection[i].waitSendMsg
	waitSendMsg <- message
}

func (client *Client) doRead() {
	for _, con := range client.connection {
		go func(conn *Conn) {
			reader := bufio.NewReader(conn.con)
			for {
				header := client8bPool.Get().([]byte)
				_, err := io.ReadFull(reader, header)
				if err != nil {
					if err == io.EOF {
						return
					}
					return
				}
				l := binary.BigEndian.Uint64(header[:8])

				client8bPool.Put(header)

				data := make([]byte, l)
				_, err = io.ReadFull(reader, data)
				if err != nil {
					if err == io.EOF {
						return
					}
					return
				}
				p := &msg.Packet{
					Header: l,
					Batch:  &data,
				}
				go client.handlePacket(p)
			}
		}(con)
	}
}

func (client *Client) handlePacket(packet *msg.Packet) {
	batch := *packet.Batch
	i := 0
	for i < len(batch) {
		l := binary.BigEndian.Uint64(batch[i : 8+i])
		seqNo := binary.BigEndian.Uint64(batch[8+i : 16+i])
		endIdx := 16 + i + int(l) - 8
		req := &msg.Msg{
			Length: l,
			SeqNo:  seqNo,
			Data:   batch[16+i : endIdx],
		}
		i = endIdx
		client.asyncHandelResponse(req)
	}
}

func (client *Client) asyncHandelResponse(ms *msg.Msg) {
	value, _ := client.ack.Load(ms.SeqNo)
	repChan, ok := value.(chan *msg.Msg)
	if !ok {
		fmt.Println("+++++", ms.SeqNo, *ms)
	}
	repChan <- ms
}

func (client *Client) doSend2() {
	for _, con := range client.connection {
		go func(conn *Conn) {
			ticker := time.NewTicker(50 * time.Millisecond)
			packet := pool.Get().([]byte)
			packetSize := 0
			for {
				select {
				case ms := <-conn.waitSendMsg:
					d := ms.Encode()
					packetSize++
					packet = append(packet, d...)

					if packetSize > BatchSize2 {
						p := msg.NewPacket(&packet)
						//io.Copy(conn.con, bytes.NewReader(p.Encode()))
						conn.con.Write(p.Encode())
						fmt.Println("+++++++++++++", packetSize)
						packet = make([]byte, 0, 10240)
						pool.Put(packet)
						packetSize = 0
					}

				case <-ticker.C:
					if len(packet) != 0 {
						p := msg.NewPacket(&packet)
						//io.Copy(conn.con, bytes.NewReader(p.Encode()))
						conn.con.Write(p.Encode())
						packet = make([]byte, 0, 10240)
						pool.Put(packet)
						packetSize = 0
					}
				}
			}
		}(con)
	}
}
